# Copyright (c) 2025 The HuggingFace Team.
# Copyright (c) 2025 Bytedance Ltd. and/or its affiliates
# SPDX-License-Identifier: Apache-2.0 
#
# This file has been modified by Bytedance Ltd. and/or its affiliates on September 15, 2025.
#
# Original file was released under Apache License 2.0, with the full license text
# available at https://github.com/huggingface/finetrainers/blob/main/LICENSE.
#
# This modified file is released under the same license.

from typing import List, Union

import torch

import finetrainers.functional as FF

from .base import ProcessorMixin


class CaptionTextDropoutProcessor(ProcessorMixin):
    def __init__(self, dropout_p: float = 0.0) -> None:
        self.dropout_p = dropout_p

    def forward(self, caption: Union[str, List[str]]) -> Union[str, List[str]]:
        return FF.dropout_caption(caption, self.dropout_p)


class CaptionEmbeddingDropoutProcessor(ProcessorMixin):
    def __init__(self, dropout_p: float = 0.0) -> None:
        self.dropout_p = dropout_p

    def forward(self, embedding: torch.Tensor) -> torch.Tensor:
        return FF.dropout_embeddings_to_zero(embedding, self.dropout_p)
